#include <stdlib.h>
#include <stdio.h>
#include "nbody.h"
#include "video.h"
#include <time.h>
#include <math.h>
#include <string.h>
#include <sys/time.h>
#include "mpi.h"
#include <mpe.h>

double *x, *y; 
double *fx, *fy, *vx, *vy, *m;
int    nbodies;
int    width;
int    height;
int    nsteps;
double g;
double interval;
double min_dist;

int    **grid;
int *displs;
int first_row, first_column, noperations, first_body, last_body;
int rango,numProcesos;
int tipo_balanceo;

/* Variables para el logging con MPE */
int start_comp, stop_comp;

/* Lee los parametros y los valores iniciales del fichero de entrda */
int init_bodies(char *file) {
	FILE *f; 
	int i;
	
	f = fopen(file,"r");
	
	fscanf(f,"%d %lf %lf %d %lf %d %d\n",&nbodies,&g,&interval,&nsteps,&min_dist,&width,&height);
	x  = malloc(sizeof(double) * nbodies);
	y  = malloc(sizeof(double) * nbodies);
	fx = malloc(sizeof(double) * nbodies);
	fy = malloc(sizeof(double) * nbodies);
	vx = malloc(sizeof(double) * nbodies);
	vy = malloc(sizeof(double) * nbodies);
	m  = malloc(sizeof(double) * nbodies);

        // creamos el vector con desplazamientos
        displs = malloc(sizeof(int)*numProcesos);
        bzero(displs, sizeof(int)*numProcesos);
	
	for (i = 0; i < nbodies; i++) {
		fscanf(f,"%lf %lf %lf %lf %lf\n",&x[i],&y[i],&vx[i],&vy[i],&m[i]);
	}
	
	fclose(f);
	
	if(rango == 0) {
		printf ("NBODY PROBLEM____________________________________________________________________\n");
		printf ("nbodies: %d\tg var: %0.2lf\tinterval: %0.2lfsec\tnsteps: %d\tmin_dist: %0.2lfp\twidth %d\theight: %d\n", nbodies,g,interval,nsteps,min_dist,width,height);
	}
	
	return 1;
}

/* Escribe los resultados al fichero de salida */
int finalize_nbodies(char *file) {
	FILE *f;
	int i;
	
	f = fopen(file,"w");
	fprintf(f,"%d %lf %lf %d %lf %d %d\n",nbodies,g,interval, nsteps,min_dist,width,height);
	for (i = 0; i < nbodies; i++) {
		fprintf(f,"%lf %lf %lf %lf %lf\n",x[i],y[i],vx[i],vy[i],m[i]);
	}
	fclose(f);

	return 1;
}

double factorial(double x) {
	double fac = 1;
	int i;
	for (i=2; i<=x; i++) fac *= i;
	return fac;
}

double mypow(double x, double exp) {
	double result = 1;
	int i;
	
	result = 1;
	for (i=0; i<exp; i++) result *= x;
	return result;
}

double mysin(double alpha) {
	double result = 0;
	int i;
	
	for (i=0; i < MAX_ACCURACY;i++) {
		result += (double)(mypow(-1,i) / factorial((double)(2*i+1))) * (double)mypow(alpha, 2*i+1);
	}
	return result;
}

double mycos(double alpha) {
	double result = 0;
	int i;
	
	for (i=0; i < MAX_ACCURACY;i++) {
		result += (double)(mypow(-1,i) / factorial((double)(2*i))) * (double)mypow(alpha, 2*i);
	}
	return result;
}

/* Decide que operaciones se ejecutarán en cada proceso para distribuir la
   carga de trabajo. Solo se usa en el caso de cargas balanceadas */
void assign_load() {
	
	int remainder, quotient, procesos, i, j, total_operations, aux;
	total_operations = 0;
	// Calculates the total amount of operations to compute the forces
	for (i = 1; i < nbodies; i++) {
		total_operations += i;
	}
	// Calculates the amount of operations in each processor
	remainder = total_operations % numProcesos;
	quotient = total_operations / numProcesos;
	procesos = 0;
	aux = 0;
	// Calculates which operations will execute each processor
	for(i = 0; i < nbodies; i++) {
		for(j = i + 1; j < nbodies; j++) {
			if (procesos < remainder) {
				noperations = quotient + 1;
				if (aux == 0) {
					if (rango == procesos){
						first_row = i;
						first_column = j;
						return;
					}
				}
				aux++;
				if(aux >= noperations) {
					procesos++;
					aux = 0;
				}
			}else{
				noperations = quotient;
				if (aux == 0) {
					if (rango == procesos){
						first_row = i;
						first_column = j;
						return;
					}
				}
				aux++;
				if(aux >= noperations) {
					procesos++;
					aux = 0;
				}					
			}
		}
	}
}


/* Decide que cuerpos se asignan a cada proceso */	
void assign_bodies() {
	int remainder, quotient, procesos, aux;
	remainder = nbodies % numProcesos;
	quotient = nbodies / numProcesos;
	aux = 0;
        //int *tmp = malloc(sizeof(int)*numProcesos);
        //bzero(tmp, sizeof(int)*numProcesos);
	
        for (procesos = 0; procesos < numProcesos; procesos++) {
		if (procesos < remainder) {
			if (rango == procesos) {
				first_body = aux;
				last_body = aux + quotient;
                                //displs[procesos] = first_body;
                                //tmp[procesos] = first_body;
				return;
			}
			aux += quotient + 1;
		} else {
			if (rango == procesos) {
				first_body = aux;
				last_body = aux + quotient - 1;
                                //displs[procesos] = first_body;
                                //tmp[procesos] = first_body;
				return;
			}
			aux += quotient;
		}
                
                //displs[procesos] = aux;
	}
        displs[0] = 0;
        displs[1] = 10;
        displs[2] = 20;
        displs[3] = 30;

	//MPI_Allgather(displs+rango, 1, MPI_INT, displs, 4, MPI_INT, MPI_COMM_WORLD);
	//MPI_Reduce (tmp, displs, numProcesos, MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);
	//MPI_Bcast(displs, numProcesos, MPI_INT, 0, MPI_COMM_WORLD);
}

/* It calculates the forces exerted on each body by the others as well as the new velocity and position of them (only used in the load balanced version)*/
int calc_nbodies_load() {
	double *ax, *ay, *fxaux, *fyaux, *vxaux, *vyaux, *xaux, *yaux;
	double fax, fay;
	double dist, f, alpha;
        //double aux;
	int i,j, count_operations;
	
	
	/* Start the calculation of the computation time with MPE */
	MPE_Log_event(start_comp,0,"start computation (loads balanced)");
	
	ax = malloc(sizeof(double) * nbodies);
	ay = malloc(sizeof(double) * nbodies);
	fxaux = malloc(sizeof(double) * nbodies);
	fyaux = malloc(sizeof(double) * nbodies);
	vxaux = malloc(sizeof(double) * nbodies);
	vyaux = malloc(sizeof(double) * nbodies);
	xaux = malloc(sizeof(double) * nbodies);
	yaux = malloc(sizeof(double) * nbodies);
	
	for (i = 0; i < nbodies; i++) {
		fxaux[i] = fyaux[i] = ax[i] = ay[i] = vxaux[i] = vyaux[i] = xaux[i] = yaux[i] = 0;
	}
	count_operations = 0;
	j = first_column;
	// Computes the forces (each processor will calculate only the operations it has assigned)
	for (i = first_row; (i < nbodies) && (count_operations < noperations); i++) {
		for (; (j < nbodies) && (count_operations < noperations); j++) {
			dist = sqrt((x[j] - x[i]) * (x[j] - x[i]) + (y[j] - y[i]) * (y[j] - y[i]));
			if (dist > min_dist ) {
				f = g * m[i] * m[j] / (dist * dist);
				if (f > MAX_F) {
					f = MAX_F;
				}

		                if (y[j]-y[i] == 0)
               		            alpha = atan(0);
		                else if (x[j]-x[i] == 0)
			            if (y[j]-y[i] > 0)
	               		        alpha = M_PI_2;
			            else
	               		        alpha = - M_PI_2;
		                else
               		            alpha = atan(y[j]-y[i] / x[j]-x[i]);

				//aux  = (y[j]-y[i]) / (x[j]-x[i]);
				//alpha = atan(aux);
				fax   = f * mycos(alpha);
				fay   = f * mysin(alpha);
				fxaux[i] += fax;
				fyaux[i] += fay;
				fxaux[j] -= fax; 
				fyaux[j] -= fay;
			}
			count_operations++;
		}
		j = i + 2;
	}
	
	// Sends the results to the processor with rank = 0 and sums up the values to compute the total forces
	MPI_Reduce (fxaux, fx, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (fyaux, fy, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	
	// Sends the total forces from the processor with rank = 0 to all the processors
	MPI_Bcast(fx, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(fy, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	
	// Calculates the new velocity and position of each body (each processor will calculate only the velocities and positions of the bodies it has assigned)
	for (i = first_body; i <= last_body; i++) {
		ax[i] = fx[i] / m[i];
		ay[i] = fy[i] / m[i];
		vxaux[i] = vx[i] + ax[i] * interval;
		vyaux[i] = vy[i] + ay[i] * interval;
		xaux[i] = x[i] + vxaux[i] * interval;
		yaux[i] = y[i] + vyaux[i] * interval;
		
		if (xaux[i] <= 0) {
			xaux[i] = 2;
			vxaux[i] = - vxaux[i];
		}
		if (xaux[i] >= width ) {
			xaux[i] = width - 2;
			vxaux[i] = - vxaux[i];
		}
		if (yaux[i] <= 0) {
			yaux[i] = 2;
			vyaux[i] = - vyaux[i];
		}
		if (yaux[i] >= height) {
			yaux[i] = height - 2;
			vyaux[i] = - vyaux[i];
		}
	}
	
	// Sends the new velocities and positions from each processor to the process with rank = 0 which will gather all the information
	MPI_Reduce (vxaux, vx, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (vyaux, vy, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (xaux, x, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (yaux, y, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	
	// The processor with rank = 0 sends all the new velocities and positions to all the processors
	MPI_Bcast(vx, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(vy, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(x, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(y, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	
	free(ax);
	free(ay);
	free(fxaux);
	free(fyaux);
	free(vxaux);
	free(vyaux);
	free(xaux);
	free(yaux);
	
	/* Stop counting time for the computation calculation */
	MPE_Log_event(stop_comp,0,"stop computation (bodies balanced)");
	
	return 1;
}

/* It calculates the forces exerted on each body by the others as well as the new velocity and position of them (only used in the body balanced version)*/
int calc_nbodies_bodies() {
	double *ax, *ay, *fxaux, *fyaux, *vxaux, *vyaux, *xaux, *yaux;
	double fax, fay;
	double dist, f, alpha;
        //double aux;
	int i,j;
        int numBodies = (last_body-first_body)+1;
        int *recv_counts = malloc(sizeof(int) * numProcesos);
        for(i = 0; i < numProcesos; i++)
            recv_counts[i] = numBodies;
	
	/* Start the calculation of the computation time with MPE */
	MPE_Log_event(start_comp,0,"start computation (bodies balanced)");
	
	ax = malloc(sizeof(double) * nbodies);
	ay = malloc(sizeof(double) * nbodies);
	fxaux = malloc(sizeof(double) * nbodies);
	fyaux = malloc(sizeof(double) * nbodies);
	vxaux = malloc(sizeof(double) * nbodies);
	vyaux = malloc(sizeof(double) * nbodies);
	xaux = malloc(sizeof(double) * nbodies);
	yaux = malloc(sizeof(double) * nbodies);
	
	for (i = 0; i < nbodies; i++) {
		fxaux[i] = fyaux[i] = ax[i] = ay[i] = vxaux[i] = vyaux[i] = xaux[i] = yaux[i] = 0;
	}
	
	// Computes the forces (each processor will calculate only the operations of the bodies it has assigned)
	for (i = first_body; i <= last_body; i++) {
		for (j = i + 1; j < nbodies; j++) {
			dist = sqrt((x[j] - x[i]) * (x[j] - x[i]) + (y[j] - y[i]) * (y[j] - y[i]));
			if (dist > min_dist ) {
				f = g * m[i] * m[j] / (dist * dist);
				if (f > MAX_F) {
					f = MAX_F;
				}

		                if (y[j]-y[i] == 0)
               		            alpha = atan(0);
		                else if (x[j]-x[i] == 0)
			            if (y[j]-y[i] > 0)
	               		        alpha = M_PI_2;
			            else
	               		        alpha = - M_PI_2;
		                else
               		            alpha = atan(y[j]-y[i] / x[j]-x[i]);

				//aux  = (y[j]-y[i]) / (x[j]-x[i]);
				//alpha = atan(aux);
				fax   = f * mycos(alpha);
				fay   = f * mysin(alpha);
				fxaux[i] += fax;
				fyaux[i] += fay;
				fxaux[j] -= fax; 
				fyaux[j] -= fay;
			}
		}
	}
	
	// Sends the results to the processor with rank = 0 and sums up the values to compute the total forces
	MPI_Reduce (fxaux, fx, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	MPI_Reduce (fyaux, fy, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	
	// Sends the total forces from the processor with rank = 0 to all the processors
	MPI_Bcast(fx, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	MPI_Bcast(fy, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	
	// Calculates the new velocity and position of each body (each processor will calculate only the velocities and positions of the bodies it has assigned)
	for (i = first_body; i <= last_body; i++) {
		ax[i] = fx[i] / m[i];
		ay[i] = fy[i] / m[i];
		vxaux[i] = vx[i] + ax[i] * interval;
		vyaux[i] = vy[i] + ay[i] * interval;
		xaux[i] = x[i] + vxaux[i] * interval;
		yaux[i] = y[i] + vyaux[i] * interval;
		
		if (xaux[i] <= 0) {
			xaux[i] = 2;
			vxaux[i] = - vxaux[i];
		}
		if (xaux[i] >= width ) {
			xaux[i] = width - 2;
			vxaux[i] = - vxaux[i];
		}
		if (yaux[i] <= 0) {
			yaux[i] = 2;
			vyaux[i] = - vyaux[i];
		}
		if (yaux[i] >= height) {
			yaux[i] = height - 2;
			vyaux[i] = - vyaux[i];
		}
	}

        displs[0] = 0;
        displs[1] = 10;
        displs[2] = 20;
        displs[3] = 30;
	
	// Sends the new velocities and positions from each processor to the process with rank = 0 which will gather all the information
	//MPI_Reduce (vxaux, vx, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	//MPI_Reduce (vyaux, vy, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	//MPI_Reduce (xaux, x, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	//MPI_Reduce (yaux, y, nbodies, MPI_DOUBLE, MPI_SUM, 0, MPI_COMM_WORLD);
	
	// The processor with rank = 0 sends all the new velocities and positions to all the processors
	//MPI_Bcast(vx, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	//MPI_Allgather(vxaux, nbodies, MPI_DOUBLE, vx, nbodies, MPI_DOUBLE, MPI_COMM_WORLD);
        //printf("Rango:%d  size:%d  pos inicial:%d\n", rango, numBodies, displs[rango]);
        /*for(i=0; i<width; i++) {
            for(j = 0; j<height; j++)
                printf("%d ", (i*width)+j);
            
            printf("\n");
        }*/

	MPI_Allgatherv(vxaux + first_body, numBodies, MPI_DOUBLE, vx, recv_counts, displs, MPI_DOUBLE, MPI_COMM_WORLD);

	//MPI_Bcast(vy, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	//MPI_Allgather(vyaux, nbodies, MPI_DOUBLE, vy, nbodies, MPI_DOUBLE, MPI_COMM_WORLD);
	MPI_Allgatherv(vyaux + first_body, numBodies, MPI_DOUBLE, vy, recv_counts, displs, MPI_DOUBLE, MPI_COMM_WORLD);

	
        //MPI_Bcast(x, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	//MPI_Allgather(xaux, nbodies, MPI_DOUBLE, x, nbodies, MPI_DOUBLE, MPI_COMM_WORLD);
	MPI_Allgatherv(xaux + first_body, numBodies, MPI_DOUBLE, x, recv_counts, displs, MPI_DOUBLE, MPI_COMM_WORLD);
	
        
        //MPI_Bcast(y, nbodies, MPI_DOUBLE, 0, MPI_COMM_WORLD);
	//MPI_Allgather(yaux, nbodies, MPI_DOUBLE, y, nbodies, MPI_DOUBLE, MPI_COMM_WORLD);
	MPI_Allgatherv(yaux + first_body, numBodies, MPI_DOUBLE, y, recv_counts, displs, MPI_DOUBLE, MPI_COMM_WORLD);
	
		
	free(ax);
	free(ay);
	free(fxaux);
	free(fyaux);
	free(vxaux);
	free(vyaux);
	free(xaux);
	free(yaux);
	
	/* Stop counting time for the computation calculation */
	MPE_Log_event(stop_comp,0,"stop computation (bodies balanced)");
	
	return 1;
}

int main (int argc, char ** argv) {
	int i, iter;
	int x11 = 0;
	double tstart, tfinish, totalTime;
	iter=0;
	
	if (argc < 5) {
		printf ("Usage:\n");
		printf ("   %s <input_file> <output_file> <X11_0|1> <assignment_strategy>\n", argv[0]);
		printf ("Where:\n");
		printf ("   <X11_0|1> can be 0 or 1. 0 means no X11, 1 means use X11.\n");
		printf ("   <assignment_strategy> can be 0 or 1. 0 means even balance of bodies, 1 means even balance of workload.\n");
		return 0;
	}
	if (!strcmp(argv[3],"1")){
		x11 = 1;
	}
	
	if(strcmp(argv[4],"1") == 0) {
		tipo_balanceo = 1;
	} else {
		if(strcmp(argv[4],"0") == 0) {
			tipo_balanceo = 0;
		} else {
			printf ("Usage:\n");
			printf ("   %s <input_file> <output_file> <X11_0|1> <assignment_strategy>\n", argv[0]);
			printf ("Where:\n");
			printf ("   <X11_0|1> can be 0 or 1. 0 means no X11, 1 means use X11.\n");
			printf ("   <assignment_strategy> can be 0 or 1. 0 means even balance of bodies, 1 means even balance of workload.\n");
			return 0;
		}
	}
	/* Initialize MPI */
	MPI_Init(&argc, &argv);
	/* Initialize MPE logging */
	MPE_Init_log();
	start_comp = MPE_Log_get_event_number();
	stop_comp = MPE_Log_get_event_number();
	MPE_Describe_state(start_comp,stop_comp,"Computation", "red");
	/* Find out the number of processes */
	MPI_Comm_size(MPI_COMM_WORLD, &numProcesos);
	/* Find out my identity in the default communicator */
	MPI_Comm_rank(MPI_COMM_WORLD, &rango);
	/* All the processors read the parameters and the initial values from the input file */
	init_bodies(argv[1]);
	/* Start the calculation of the computation time with MPE */
	MPE_Log_event(start_comp,0,"start computation (bodies balanced)");
	// The process with rank = 0 sets up the window
	if (x11 && rango == 0) {
		allocate_grid(width, height);
		setupWindow(width, height);
	}
	MPE_Log_event(stop_comp,0,"stop computation (bodies balanced)");
	// In the load balanced version, each processor calculates which operations will have to compute
	if (tipo_balanceo == 1) {
		assign_load();
	}
	// Each processor calculates which bodies have assigned
	assign_bodies();
	
	// Each processor executes the body of this loop once on each step of the program
	totalTime = 0;
	for (iter = 0; iter < nsteps; iter++) {
		tstart = MPI_Wtime();
		// The processors compute the forces, velocities and positions of the bodies
		if (tipo_balanceo == 1) {
			calc_nbodies_load();
		} else {
			calc_nbodies_bodies();
		}
		tfinish = MPI_Wtime();
		totalTime += tfinish - tstart;
		// The process with rank = 0 updates the position of the bodies in the window
		if (rango == 0) {
			if (x11) {
				MPE_Log_event(start_comp,0,"start computation (bodies balanced)");
				init_grid(width, height);
				for (i=0; i < nbodies; i++) {
					grid[(int)x[i]][(int)y[i]]= m[i];
				}
				if (iter % REFRESH ==0 ){
					myDraw(width, height);
				}
				MPE_Log_event(stop_comp,0,"stop computation (bodies balanced)");
			}
		}
	}
	
	
	if (rango == 0) {
		if (x11) {
			cleanup_grid(width, height);
		}
		finalize_nbodies(argv[2]);
	}
	printf ("Process rank: %d   nbody takes: %fsec    mean: %fsec    accuracy: %d \n", rango, totalTime, totalTime/nsteps ,MAX_ACCURACY);
	
	free(x);
	free(y);
	free(fx);
	free(fy);
	free(vx);
	free(vy);
	free(m);
        free(displs);

	MPE_Finish_log("nbody_mpi.log");
	MPI_Finalize();
	return 1;
}

